import networkx as nx
import random
import csv

# NULL  Bob Ted
# Kate  5,4 6,7
# Mary  8,9 10,11
# etc.

def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
    n = start
    while True:
        yield n
        n += step

def fromMatrixCSV(filename):
    """
    Reads in the preference matrix.
    """
    source = csv.reader(open(filename))
    leads = source.next()[1:] # Skip the blank space
    graph = nx.Graph()
    odds = count(1,2)
    evens = count(0,2)
    nameToNumber = {}
    for lead in leads:
        nameToNumber[lead] = odds.next()
        graph.add_node(nameToNumber[lead], name=lead)

    for line in source:
        nameToNumber[line[0]] = evens.next()
        graph.add_node(nameToNumber[line[0]], name=line[0])
        prefs = line[1:]
        for i in range(len(prefs)):
            pref = [int(j) for j in prefs[i].split(',')]
            graph.add_edge(nameToNumber[line[0]], nameToNumber[leads[i]],
                    weight=sum(pref), lead=pref[0], follow=pref[1])

    return graph, nameToNumber

def readAttributes(filename, graph, people):
    """
    Reads in the attribute list, combines it with the preference matrix.
    """
    source = csv.reader(open(filename))
    for person in source:
        try: pin = people[person[0]]
        except KeyError:
            print "%s not in preference list, skipping" % person[0]
            continue
        graph.node[pin]['height'] = float(person[1])
        graph.node[pin]['seniority'] = float(person[2])
        graph.node[pin]['sadness'] = float(person[3])
        graph.node[pin]['attendance'] = float(person[4])
        #graph.node[pin]['untried'] = sorted(graph[pin].keys(), cmp=(lambda x,y: cmp(graph[pin][x], graph[pin][y])))
        graph.node[pin]['untried'] = graph[pin].keys()
        random.shuffle(graph.node[pin]['untried'])

    return graph

def parseData(prefs, info):
    """
    Generates the top-level graph.
    """
    graph, people = fromMatrixCSV(prefs)
    graph = readAttributes(info, graph, people)
    return graph 
